import os
import ctypes
from typing import List, Tuple, Union, TypeVar, Generic

NAME_LENGTH = 24
# from typing import Self

dll_path = os.path.join(os.path.dirname(__file__), "lib", "core.dll")

T = TypeVar("T")


def is_null_pointer(p: Union["POINTER_TYPE[T]"]) -> bool:
    try:
        if p.contents:
            pass
        return False
    except ValueError:
        return True


class POINTER_TYPE(Generic[T]):
    @property
    def contents(self) -> T:
        raise ValueError(
            "Do not actually call this method because this class is a type hint"
        )

    def __getitem__(self, item: int) -> T:
        raise ValueError()


# 假设的枚举类型定义，这需要根据实际情况来调整
class ELECTRIC_MOTOR_TYPE(ctypes.c_int32):
    THREE_PHASED_ASYNC = 0
    THREE_PHASED_SYNC = 1
    SINGLE_PHASED = 2
    DC_BRUSHED = 3


class LOCOMOTIVE_TYPE(ctypes.c_int32):
    LOCOMOTIVE_TYPE_ELECTRIC = 0
    LOCOMOTIVE_TYPE_DIESEL_ELECTRIC = 1
    LOCOMOTIVE_TYPE_STEAM = 2


class RAILROAD_CAR_TYPE(ctypes.c_int32):
    RAILROAD_CAR_TYPE_FLATCAR = 0


class TRAIN_STATE(ctypes.c_int32):
    TRAIN_STATE_STOP = 0
    TRAIN_STATE_ACCEL = 1
    TRAIN_STATE_BRAKING = 2


class SIGNAL(ctypes.c_int32):
    TRAIN_STOP = 0
    TRAIN_WARNING = 1
    TRAIN_PASS = 2


class DIRECTION(ctypes.c_int32):
    DIRECTION_DOWN = 0
    DIRECTION_UP = 1


# 加载DLL
dll = ctypes.CDLL(dll_path)


# 定义ElectricMotorProperty结构体
class ElectricMotorProperty(ctypes.Structure):
    _fields_ = [
        ("name", ctypes.c_char * NAME_LENGTH),
        ("maxPower", ctypes.c_double),
        ("maxTorque", ctypes.c_double),
        ("maxRPM", ctypes.c_double),
        ("efficiency", ctypes.c_double),
        ("gearRatio", ctypes.c_double),
        ("type", ELECTRIC_MOTOR_TYPE),
    ]
    name: bytes
    maxPower: float
    maxTorque: float
    maxRPM: float
    efficiency: float
    gearRatio: float
    type: ELECTRIC_MOTOR_TYPE

    @classmethod
    def create_value(
        cls,
        name: str,
        type: ELECTRIC_MOTOR_TYPE,
        maxPower: float,
        maxTorque: float,
        maxRPM: float,
        efficiency: float,
        gearRatio: float,
    ):
        return dll.ElectricMotorProperty_createValue(
            name, type, maxPower, maxTorque, maxRPM, efficiency, gearRatio
        )


class ElectricMotor(ctypes.Structure):
    _fields_ = [("property", ElectricMotorProperty), ("status", ctypes.c_char)]

    @classmethod
    def create_multi(
        cls,
        motor_num: int,
    ) -> ctypes.POINTER:
        return dll.ElectricMotor_createMulti(
            motor_num,
        )


class Locomotive(ctypes.Structure):
    _fields_ = [
        ("type", LOCOMOTIVE_TYPE),  # 机车类型
        ("maxPower", ctypes.c_double),  # 最大功率 kW [CACHE]
        ("startSpeed", ctypes.c_double),  # 具备启动牵引力的最大速度 km/h
        ("mass", ctypes.c_double),  # 总质量 t
        ("adhesionMass", ctypes.c_double),  # 黏着重量 t
        ("wheelDiameter", ctypes.c_double),  # 驱动轮直径 m
        ("axleLoad", ctypes.c_double),  # 轴重 t
    ]


class ElectricLocomotive(ctypes.Structure):
    _fields_ = [
        ("base", Locomotive),  # 机车类型
        ("motorProperty", ElectricMotorProperty),  # 牵引电机属性
        ("motorNum", ctypes.c_int32),  # 最大功率 kW [CACHE]
        ("motors", ctypes.POINTER(ElectricMotor)),  # 具备启动牵引力的最大速度 km/h
    ]
    base: Locomotive
    motorProperty: ElectricMotorProperty
    motorNum: int
    motors: "ctypes._Pointer[ElectricMotor]"

    @classmethod
    def create(
        cls,
        startSpeed: float,
        axleLoad: float,
        mass: float,
        wheelDiameter: float,
        motorNum: int,
        motorProperty: "ElectricMotorProperty",
    ) -> Tuple[ctypes.POINTER, "ElectricLocomotive"]:
        ptr = dll.Locomotive_create_electric(
            startSpeed, axleLoad, mass, wheelDiameter, motorNum, motorProperty
        )
        return ptr, ptr.contents

    def traction_limit_torque(self) -> float:
        return dll.Locomotive_tractionLimitTorque(
            ctypes.cast(ctypes.pointer(self), ctypes.POINTER(Locomotive))
        )

    def traction_limit_adhesion(self, speed: float) -> float:
        return dll.Locomotive_tractionLimitAdhesion(
            ctypes.cast(ctypes.pointer(self), ctypes.POINTER(Locomotive)), speed
        )

    def speed_limit_rpm(self) -> float:
        return dll.Locomotive_speedLimitRPM(
            ctypes.cast(ctypes.pointer(self), ctypes.POINTER(Locomotive))
        )

    def traction_limit_power(self, power: float, v: float) -> float:
        return dll.Locomotive_tractionLimitPower(
            ctypes.cast(ctypes.pointer(self), ctypes.POINTER(Locomotive)), power, v
        )

    def update_parameters(self):
        dll.ElectricLocomotive_updateParameters(ctypes.pointer(self))


class RailroadCarProperty(ctypes.Structure):
    _fields_ = [
        ("name", ctypes.c_char * NAME_LENGTH),
        ("type", RAILROAD_CAR_TYPE),
        ("length", ctypes.c_double),
        ("designSpeed", ctypes.c_double),
        ("deadWeight", ctypes.c_double),
        ("maxLoad", ctypes.c_double),
    ]
    name: str
    type: RAILROAD_CAR_TYPE
    length: float
    designSpeed: float
    deadWeight: float
    maxLoad: float

    @classmethod
    def create(
        cls,
        name: bytes,
        type: RAILROAD_CAR_TYPE,
        length: float,
        designSpeed: float,
        deadWeight: float,
        maxLoad: float,
    ):
        name_ = name
        return dll.RailroadCarProperty_create(
            name_,
            type,
            length,
            designSpeed,
            deadWeight,
            maxLoad,
        ).contents


class RailroadWagon(ctypes.Structure):
    _fields_ = [
        ("id", ctypes.c_int32),
        ("property", ctypes.POINTER(RailroadCarProperty)),
        ("load", ctypes.c_double),
    ]
    id: int
    property: "ctypes._Pointer[RailroadCarProperty]"
    load: float

    @classmethod
    def create(cls, property: RailroadCarProperty) -> "RailroadWagon":
        return dll.RailroadWagon_create(ctypes.pointer(property)).contents


class TrainDynamicProperties(ctypes.Structure):
    _fields_ = [
        ("speed", ctypes.c_double),  # 速度，单位km/h
        ("acceleration", ctypes.c_double),  # 加速度，单位m/s^2
        ("mass", ctypes.c_double),  # 质量，单位t
        ("distance", ctypes.c_double),  # 距离，单位m
        ("targetSpeed", ctypes.c_double),  # 目标速度，单位km/h
        ("direction", ctypes.c_double),  # 0:正向，1:反向
        ("length", ctypes.c_double),  # 火车长度（米）
        ("state", TRAIN_STATE),
    ]
    speed: float
    acceleration: float
    mass: float
    distance: float
    targetSpeed: float
    direction: float
    length: float
    state: TRAIN_STATE


class Track(ctypes.Structure):
    _fields_ = [
        ("distance", ctypes.c_double),  # 区间长度，单位km
        ("head_signal", SIGNAL),  # 上行方向信号（位于区间0位置）
        ("tail_signal", SIGNAL),  # 下行方向信号（位于区间末尾位置）
        ("trains", ctypes.c_void_p),  # 区间内列车列表（链表指针，不要通过Python修改）
        ("id", ctypes.c_int32),  # 区间ID
    ]
    distance: float
    trains: ctypes.c_void_p
    id: int

    @classmethod
    def create(cls, id: int) -> "Track":
        return dll.Track_create(id).contents

    def add_train(self, train: "Train"):
        dll.Track_addTrain(ctypes.pointer(self), ctypes.pointer(train))

    def update(self):
        dll.Track_update(ctypes.pointer(self))


class RouteItem(ctypes.Structure):
    _fields_ = [("track", ctypes.POINTER(Track)), ("direction", ctypes.c_int32)]
    track: POINTER_TYPE[Track]
    direction: DIRECTION


class Route(ctypes.Structure):
    _fields_ = [
        ("items", ctypes.POINTER(RouteItem)),
        ("length", ctypes.c_int32),
        ("index", ctypes.c_int32),
    ]
    items: POINTER_TYPE[RouteItem]
    length: int
    index: int

    @classmethod
    def create(cls, tracks: List[Track], directions: List[DIRECTION]) -> "Route":
        assert len(tracks) == len(directions)
        tracks_ptr_array = (ctypes.POINTER(Track) * len(tracks))()
        directions_array = (ctypes.c_int32 * len(directions))()
        for i, track, direction in zip(range(len(tracks)), tracks, directions):
            tracks_ptr_array[i] = ctypes.pointer(track)
            directions_array[i] = direction

        return dll.Route_create(
            ctypes.cast(
                ctypes.pointer(tracks_ptr_array), ctypes.POINTER(ctypes.POINTER(Track))
            ),
            ctypes.cast(ctypes.pointer(directions_array), ctypes.POINTER(DIRECTION)),
            len(tracks),
        ).contents

    def destroy(self):
        dll.Route_destroy(ctypes.pointer(self))


class Train(ctypes.Structure):
    _fields_ = [
        ("locomotives", ctypes.POINTER(ctypes.POINTER(Locomotive))),
        ("locomotivesNum", ctypes.c_int32),
        ("wagons", ctypes.POINTER(ctypes.POINTER(RailroadWagon))),
        ("wagonsNum", ctypes.c_int32),
        ("trainDynamics", TrainDynamicProperties),
        ("route", ctypes.POINTER(Route)),  # ctypes绝对不允许在POINTER中加字符串！
    ]
    locomotives: "ctypes.POINTER(ctypes.POINTER(Locomotive))"
    locomotivesNum: int
    trainDynamics: TrainDynamicProperties
    route: POINTER_TYPE["Route"]

    @classmethod
    def create(
        cls, locomotives: "list[Locomotive]", wagons: "list[RailroadWagon]"
    ) -> "Train":
        loco_ptr_array = (ctypes.POINTER(Locomotive) * len(locomotives))()
        for i in range(len(locomotives)):
            loco_ptr_array[i] = ctypes.cast(
                ctypes.pointer(locomotives[i]), ctypes.POINTER(Locomotive)
            )
        wagons_ptr_array = (ctypes.POINTER(RailroadWagon) * len(wagons))()
        for i, wagon in enumerate(wagons):
            wagons_ptr_array[i] = ctypes.cast(
                ctypes.pointer(wagon), ctypes.POINTER(RailroadWagon)
            )
        return dll.Train_create(
            ctypes.cast(loco_ptr_array, ctypes.POINTER(ctypes.POINTER(Locomotive))),
            len(locomotives),
            ctypes.cast(
                wagons_ptr_array, ctypes.POINTER(ctypes.POINTER(RailroadWagon))
            ),
            len(wagons),
        ).contents

    def set_route(self, route: "Route") -> None:
        return dll.Train_setRoute(ctypes.pointer(self), ctypes.pointer(route))

    def update(
        self,
    ) -> None:
        return dll.Train_update(ctypes.pointer(self), None)

    def get_position(self) -> Tuple[int,float]:
        """
        获取行进了当前区间的百分之多少
        """
        if self.route.contents.index >= self.route.contents.length:
            return -1, -1.0
        current_track = self.route.contents.items[
            self.route.contents.index
        ].track.contents
        return current_track.id, self.trainDynamics.distance / current_track.distance


def set_func_signature(func, argtypes: list, restype):
    func.argtypes = argtypes
    func.restype = restype


def _setup_dll_types():
    set_func_signature(
        dll.ElectricMotorProperty_createValue,
        [
            ctypes.c_char_p,
            ELECTRIC_MOTOR_TYPE,
            ctypes.c_double,
            ctypes.c_double,
            ctypes.c_double,
            ctypes.c_double,
            ctypes.c_double,
        ],
        ElectricMotorProperty,
    )
    set_func_signature(
        dll.ElectricMotor_createMulti,
        [
            ctypes.c_int32,
        ],
        ctypes.POINTER(ElectricMotor),
    )

    dll.Locomotive_create_electric.argtypes = [
        ctypes.c_double,
        ctypes.c_double,
        ctypes.c_double,
        ctypes.c_double,
        ctypes.c_int32,
        ElectricMotorProperty,
    ]
    dll.Locomotive_create_electric.restype = ctypes.POINTER(ElectricLocomotive)

    dll.Locomotive_tractionLimitAdhesion.argtypes = [
        ctypes.POINTER(Locomotive),
        ctypes.c_double,
    ]
    dll.Locomotive_tractionLimitAdhesion.restype = ctypes.c_double

    set_func_signature(
        dll.RailroadCarProperty_create,
        [
            ctypes.c_char_p,
            RAILROAD_CAR_TYPE,
            ctypes.c_double,
            ctypes.c_double,
            ctypes.c_double,
            ctypes.c_double,
        ],
        ctypes.POINTER(RailroadCarProperty),
    )
    set_func_signature(
        dll.RailroadWagon_create,
        [ctypes.POINTER(RailroadCarProperty)],
        ctypes.POINTER(RailroadWagon),
    )

    set_func_signature(
        dll.ElectricLocomotive_updateParameters,
        [ctypes.POINTER(ElectricLocomotive)],
        None,
    )
    set_func_signature(
        dll.Locomotive_tractionLimitPower,
        [
            ctypes.POINTER(Locomotive),
            ctypes.c_double,
            ctypes.c_double,
        ],
        ctypes.c_double,
    )

    set_func_signature(
        dll.Locomotive_tractionLimitTorque,
        [
            ctypes.POINTER(Locomotive),
        ],
        ctypes.c_double,
    )
    set_func_signature(
        dll.Locomotive_speedLimitRPM,
        [
            ctypes.POINTER(Locomotive),
        ],
        ctypes.c_double,
    )
    set_func_signature(
        dll.Train_setRoute, [ctypes.POINTER(Train), ctypes.POINTER(Route)], None
    )
    set_func_signature(
        dll.Train_update, [ctypes.POINTER(Train), ctypes.POINTER(Train)], None
    )
    set_func_signature(
        dll.Train_create,
        [
            ctypes.POINTER(ctypes.POINTER(Locomotive)),
            ctypes.c_int32,
            ctypes.POINTER(ctypes.POINTER(RailroadWagon)),
            ctypes.c_int32,
        ],
        ctypes.POINTER(Train),
    )
    set_func_signature(
        dll.Track_create,
        [
            ctypes.c_int32,
        ],
        ctypes.POINTER(Track),
    )
    set_func_signature(
        dll.Track_addTrain,
        [ctypes.POINTER(Track), ctypes.POINTER(Train)],
        None,
    )
    set_func_signature(
        dll.Track_update,
        [ctypes.POINTER(Track)],
        None,
    )
    set_func_signature(
        dll.Route_create,
        [
            ctypes.POINTER(ctypes.POINTER(Track)),
            ctypes.POINTER(ctypes.c_int32),
            ctypes.c_int32,
        ],
        ctypes.POINTER(Route),
    )
    set_func_signature(dll.Route_destroy, [ctypes.POINTER(Route)], None)


_setup_dll_types()
